import cv2
from mmocr.utils.ocr import MMOCR


class MLT_Det:
    def __init__(self, det_config, det_checkpoint):
        self.det_model = MMOCR(det='PANet_IC15',
                                 recog=None,
                                 det_config=det_config,
                                 det_ckpt=det_checkpoint)

    def __call__(self, im):
        '''

        :param im: numpy
        :return: string
        '''
        recog_text = self.det_model.readtext(im,imshow=True)
        return recog_text

    def predict_polys(self,im):
        res=self.det_model.readtext(im)
        pred_boxes_ws = res[0]['boundary_result']
        pred_boxes = []
        for pred_box in pred_boxes_ws:
            pred_boxes.append(pred_box[:-1])
        return pred_boxes

if __name__ == '__main__':
    det_config = '/home/wsl/OCR/mmocr_mlt/configs/textdet/psenet/psenet_r50_fpnf_600e_mlt19.py'
    det_checkpoint = '/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/mlt_ckpt/det/psenet.pth'
    model = MLT_Det(det_config, det_checkpoint)
    im = cv2.imread('/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/PublicDataSet/OCR/MLT17/ch8_training_images_1/img_10.jpg')
    res = model.predict_polys(im)
    print(res)